package com.zhang.spark_1.spark_core.acc

import org.apache.spark.rdd.RDD
import org.apache.spark.util.AccumulatorV2
import org.apache.spark.{SparkConf, SparkContext}

import scala.collection.mutable

/**
 * @title:
 * @author: zhang
 * @date: 2021/12/5 16:06 
 */
object Spark03_ACC_wordcount {

  def main(args: Array[String]): Unit = {
    val conf: SparkConf = new SparkConf().setMaster("local[*]").setAppName("ACC")
    val sc = new SparkContext(conf)

    val rdd: RDD[String] = sc.makeRDD(List("hello","scala","hello","spark"))


    //var sum = 0;

    //使用自定义累加器
    val wcAcc = new MyAccumulator()
    sc.register(wcAcc,"wordcount")

    rdd.foreach(
      word=>{
        wcAcc.add(word)
      }
    )

    println(wcAcc.value)


    sc.stop()
  }

  class MyAccumulator extends AccumulatorV2[String, mutable.Map[String,Int]] {
    private var wcMap = mutable.Map[String,Int]()

    override def isZero: Boolean = {
      wcMap.isEmpty
    }

    override def copy(): AccumulatorV2[String, mutable.Map[String, Int]] = {
      new MyAccumulator
    }

    override def reset(): Unit = {
      wcMap.clear()
    }

    override def add(word: String): Unit = {
      val newCnt = wcMap.getOrElse(word,0)+1
      wcMap.update(word,newCnt)
    }

    override def merge(other: AccumulatorV2[String, mutable.Map[String, Int]]): Unit = {
      var map1 = this.wcMap
      var map2 = other.value
      map2.foreach{
        case (word,cont)=>{
          val newCount = map1.getOrElse(word,0)+cont
          wcMap.update(word,newCount)
        }
      }
    }

    override def value: mutable.Map[String, Int] = {
      wcMap
    }
  }
}
